"""
数据库连接管理器
负责管理数据库连接、事务处理和查询执行
"""

import logging
from typing import Dict, List, Any, Optional, Union
from contextlib import contextmanager
import threading
from sqlalchemy import create_engine, text, inspect
from sqlalchemy.orm import sessionmaker
from sqlalchemy.engine import Engine
from sqlalchemy.exc import SQLAlchemyError

from config.database_config import db_config_manager, DatabaseConfig, DatabaseType


class DatabaseConnectionManager:
    """数据库连接管理器"""
    
    def __init__(self):
        self.engines: Dict[str, Engine] = {}
        self.sessions: Dict[str, sessionmaker] = {}
        self.lock = threading.Lock()
        self.logger = logging.getLogger(__name__)
    
    def _get_connection_string(self, config: DatabaseConfig) -> str:
        """构建数据库连接字符串"""
        if config.db_type == DatabaseType.MYSQL:
            return f"mysql+pymysql://{config.username}:{config.password}@{config.host}:{config.port}/{config.database}?charset=utf8mb4"
        elif config.db_type == DatabaseType.POSTGRESQL:
            return f"postgresql://{config.username}:{config.password}@{config.host}:{config.port}/{config.database}"
        elif config.db_type == DatabaseType.SQLITE:
            return f"sqlite:///{config.database}"
        elif config.db_type == DatabaseType.SQLSERVER:
            return f"mssql+pyodbc://{config.username}:{config.password}@{config.host}:{config.port}/{config.database}?driver=ODBC+Driver+17+for+SQL+Server"
        else:
            raise ValueError(f"不支持的数据库类型: {config.db_type}")
    
    def get_engine(self, config_name: str) -> Optional[Engine]:
        """获取数据库引擎"""
        if config_name in self.engines:
            return self.engines[config_name]
        
        config = db_config_manager.get_config(config_name)
        if not config:
            self.logger.error(f"未找到数据库配置: {config_name}")
            return None
        
        try:
            connection_string = self._get_connection_string(config)
            engine = create_engine(
                connection_string,
                pool_pre_ping=True,
                pool_recycle=3600,
                echo=False
            )
            
            with self.lock:
                self.engines[config_name] = engine
                self.sessions[config_name] = sessionmaker(bind=engine)
            
            self.logger.info(f"数据库引擎创建成功：{config_name}")
            return engine
            
        except Exception as e:
            self.logger.error(f"创建数据库引擎失败 {config_name}: {str(e)}")
            return None
    
    def test_connection(self, config_name: str) -> bool:
        """测试数据库连接"""
        try:
            engine = self.get_engine(config_name)
            if not engine:
                return False
            
            with engine.connect() as conn:
                conn.execute(text("SELECT 1"))
            
            self.logger.info(f"数据库连接测试成功: {config_name}")
            return True
            
        except Exception as e:
            self.logger.error(f"数据库连接测试失败 {config_name}: {str(e)}")
            return False
    
    def get_schema_names(self, config_name: str) -> List[str]:
        """获取数据库架构名称列表"""
        try:
            engine = self.get_engine(config_name)
            if not engine:
                return []
            
            inspector = inspect(engine)
            schemas = inspector.get_schema_names()
            
            # 对于某些数据库，添加默认架构
            config = db_config_manager.get_config(config_name)
            if config:
                if config.db_type == DatabaseType.MYSQL:
                    schemas.insert(0, config.database)
                elif config.db_type == DatabaseType.SQLITE:
                    schemas = ["main"]
            
            return schemas
            
        except Exception as e:
            self.logger.error(f"获取架构名称失败 {config_name}: {str(e)}")
            return []
    
    def get_table_names(self, config_name: str, schema: Optional[str] = None) -> List[str]:
        """获取表名称列表"""
        try:
            engine = self.get_engine(config_name)
            if not engine:
                return []
            
            inspector = inspect(engine)
            tables = inspector.get_table_names(schema=schema)
            
            return sorted(tables)
            
        except Exception as e:
            self.logger.error(f"获取表名称失败 {config_name}: {str(e)}")
            return []
    
    def get_table_info(self, config_name: str, table_name: str, schema: Optional[str] = None) -> Dict[str, Any]:
        """获取表详细信息"""
        try:
            engine = self.get_engine(config_name)
            if not engine:
                return {}
            
            inspector = inspect(engine)
            
            # 获取列信息
            columns = inspector.get_columns(table_name, schema=schema)
            
            # 获取主键信息
            primary_keys = inspector.get_pk_constraint(table_name, schema=schema)
            
            # 获取外键信息
            foreign_keys = inspector.get_foreign_keys(table_name, schema=schema)
            
            # 获取索引信息
            indexes = inspector.get_indexes(table_name, schema=schema)
            
            # 获取表注释
            table_comment = ""
            try:
                table_info = inspector.get_table_comment(table_name, schema=schema)
                table_comment = table_info.get('text', '') if table_info else ''
            except:
                pass
            
            return {
                'columns': columns,
                'primary_keys': primary_keys,
                'foreign_keys': foreign_keys,
                'indexes': indexes,
                'table_comment': table_comment
            }
            
        except Exception as e:
            self.logger.error(f"获取表信息失败 {config_name}.{table_name}: {str(e)}")
            return {}
    
    def execute_query(self, config_name: str, query: str, params: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
        """执行查询并返回结果"""
        try:
            engine = self.get_engine(config_name)
            if not engine:
                return []
            
            with engine.connect() as conn:
                result = conn.execute(text(query), params or {})
                columns = result.keys()
                rows = result.fetchall()
                
                return [dict(zip(columns, row)) for row in rows]
                
        except Exception as e:
            self.logger.error(f"查询执行失败 {config_name}: {str(e)}")
            return []
    
    def fetch_one(self, config_name: str, query: str, params: Optional[Dict[str, Any]] = None) -> Optional[Dict[str, Any]]:
        """执行查询并返回单行结果"""
        try:
            engine = self.get_engine(config_name)
            if not engine:
                return None
            
            with engine.connect() as conn:
                result = conn.execute(text(query), params or {})
                row = result.fetchone()
                
                if row:
                    columns = result.keys()
                    return dict(zip(columns, row))
                return None
                
        except Exception as e:
            self.logger.error(f"查询执行失败 {config_name}: {str(e)}")
            return None
    
    @contextmanager
    def get_session(self, config_name: str):
        """获取数据库会话上下文管理器"""
        session_factory = self.sessions.get(config_name)
        if not session_factory:
            engine = self.get_engine(config_name)
            if not engine:
                raise RuntimeError(f"无法获取数据库引擎: {config_name}")
            session_factory = sessionmaker(bind=engine)
            with self.lock:
                self.sessions[config_name] = session_factory
        
        session = session_factory()
        try:
            yield session
            session.commit()
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()
    
    def close_all_connections(self):
        """关闭所有数据库连接"""
        with self.lock:
            for engine in self.engines.values():
                try:
                    engine.dispose()
                except Exception as e:
                    self.logger.error(f"关闭数据库连接失败: {str(e)}")
            
            self.engines.clear()
            self.sessions.clear()
        
        self.logger.info("所有数据库连接已关闭")


# 创建全局数据库管理器实例
db_manager = DatabaseConnectionManager()